/**
 * 
 */
package edu.umd.clip.lm.model.training;

import edu.umd.clip.lm.model.training.NewTrainer.Distributions;
import edu.umd.clip.lm.util.ProbMath;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class EntropyQuestionEstimator extends AbstractQuestionEstimator {
	private final int minCount;
	private final double minCostReduction;
	  
	public EntropyQuestionEstimator(int minCount, double minCostReduction) {
		this.minCount = minCount;
		this.minCostReduction = minCostReduction;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.AbstractQuestionEstimator#prePruningRule(edu.umd.clip.lm.model.training.AbstractTrainingData, edu.umd.clip.lm.model.training.AbstractTrainingData)
	 */
	@Override
	public boolean prePruningRule(ActiveTreeNode activeNode) {
		for(long count :activeNode.counts) {
			if (count < minCount) return false;
		}
		return true;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.AbstractQuestionEstimator#estimateQuestion(edu.umd.clip.lm.model.training.ActiveTreeNode, edu.umd.clip.lm.model.training.NewTrainer.Distributions)
	 */
	@Override
	public Result estimateQuestion(int data, ActiveTreeNode activeNode, Distributions distributions) {
/*		
		double falseEntropy = ProbMath.computeEntropy(distributions.falseDist);
		double trueEntropy = ProbMath.computeEntropy(distributions.trueDist);
		double entropy = (falseEntropy*distributions.falseCount + trueEntropy*distributions.trueCount) /
			(distributions.falseCount + distributions.trueCount);
*/		
		long totalCount = getTotalCounts()[data];
		
		double falseEntropy = ProbMath.computeEntropy(distributions.falseDist) * 
			distributions.falseDist.getTotalCount() / totalCount;
		double trueEntropy = ProbMath.computeEntropy(distributions.trueDist) *
			distributions.trueDist.getTotalCount() / totalCount;
		double entropy = falseEntropy + trueEntropy;
		
		
		if (distributions.falseDist.getTotalCount() < minCount || distributions.trueDist.getTotalCount() < minCount) {
			entropy = Double.NaN;
		}
		Result result = new Result(entropy);
		result.trueCost = trueEntropy;
		result.falseCost = falseEntropy;
		return result;
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.AbstractQuestionEstimator#isGood(edu.umd.clip.lm.model.training.AbstractQuestionEstimator.Result)
	 */
	@Override
	public boolean isGood(int data, Result result, ActiveTreeNode activeNode) {
		return activeNode.nodeCosts[data] - result.cost > minCostReduction;
	}

}
